import pygame
import math
import time
import random

# Constants
SCREEN_WIDTH = 1024
SCREEN_HEIGHT = 768
CENTER = (SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2)
HEX_RADIUS = 200
BALL_RADIUS = 30
GRAVITY = 800
AIR_FRICTION = 0.9
RESTITUTION = 0.8
FRICTION = 0.3
ANGULAR_SPEED = -2
BALL_COUNT = 20
COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255)]


class Ball:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.vx = random.uniform(-100, 100)
        self.vy = random.uniform(-50, 50)
        self.color = color


def closest_point_on_segment(A, B, P):
    ax, ay = A
    bx, by = B
    px, py = P

    abx = bx - ax
    aby = by - ay
    apx = px - ax
    apy = py - ay

    ab2 = abx * abx + aby * aby
    if ab2 == 0:
        return (ax, ay)

    t = (apx * abx + apy * aby) / ab2
    t = max(0, min(1, t))

    return (ax + t * abx, ay + t * aby)


def compute_normal(A, B, center):
    mx = (A[0] + B[0]) / 2
    my = (A[1] + B[1]) / 2
    dx = center[0] - mx
    dy = center[1] - my
    length = math.hypot(dx, dy)
    return (dx / length, dy / length) if length else (0, 0)


def handle_ball_collision(b1, b2):
    dx = b2.x - b1.x
    dy = b2.y - b1.y
    dist = math.hypot(dx, dy)

    if dist < 2 * BALL_RADIUS:
        overlap = 2 * BALL_RADIUS - dist
        if dist == 0:
            return

        nx = dx / dist
        ny = dy / dist

        # Position correction
        correction = overlap * 0.5
        b1.x -= nx * correction
        b1.y -= ny * correction
        b2.x += nx * correction
        b2.y += ny * correction

        # Relative velocity
        v_rel_x = b1.vx - b2.vx
        v_rel_y = b1.vy - b2.vy

        # Impulse calculation
        impulse = (v_rel_x * nx + v_rel_y * ny) * (1 + RESTITUTION)
        impulse /= 2  # equal mass

        # Update velocities
        b1.vx -= impulse * nx
        b1.vy -= impulse * ny
        b2.vx += impulse * nx
        b2.vy += impulse * ny


pygame.init()
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Multi-Ball Spinning Hexagon")

balls = [Ball(CENTER[0] + random.uniform(-50, 50),
              CENTER[1] + random.uniform(-50, 50),
              random.choice(COLORS)) for _ in range(BALL_COUNT)]

hex_angle = 0.0
prev_time = time.time()

running = True
clock = pygame.time.Clock()

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    current_time = time.time()
    dt = current_time - prev_time
    prev_time = current_time

    # Update hexagon rotation
    hex_angle += ANGULAR_SPEED * dt

    # Update physics for all balls
    for ball in balls:
        ball.vy += GRAVITY * dt
        ball.vx *= math.exp(-AIR_FRICTION * dt)
        ball.vy *= math.exp(-AIR_FRICTION * dt)
        ball.x += ball.vx * dt
        ball.y += ball.vy * dt

    # Calculate hexagon vertices
    vertices = []
    cx, cy = CENTER
    for i in range(6):
        angle = hex_angle + i * (2 * math.pi / 6)
        x = cx + HEX_RADIUS * math.cos(angle)
        y = cy + HEX_RADIUS * math.sin(angle)
        vertices.append((x, y))

    # Handle wall collisions
    for ball in balls:
        collided = False
        for i in range(6):
            A = vertices[i]
            B = vertices[(i + 1) % 6]
            Q = closest_point_on_segment(A, B, (ball.x, ball.y))
            dx, dy = ball.x - Q[0], ball.y - Q[1]

            if dx * dx + dy * dy <= BALL_RADIUS ** 2:
                N = compute_normal(A, B, CENTER)
                qx, qy = Q
                vw_x = -ANGULAR_SPEED * (qy - cy)
                vw_y = ANGULAR_SPEED * (qx - cx)

                rel_vx = ball.vx - vw_x
                rel_vy = ball.vy - vw_y
                vn = rel_vx * N[0] + rel_vy * N[1]

                if vn < 0:
                    vt_x = rel_vx - vn * N[0]
                    vt_y = rel_vy - vn * N[1]

                    vn_new = -vn * RESTITUTION
                    vt_x *= (1 - FRICTION)
                    vt_y *= (1 - FRICTION)

                    ball.vx = vw_x + vn_new * N[0] + vt_x
                    ball.vy = vw_y + vn_new * N[1] + vt_y
                    ball.x, ball.y = Q[0] + N[0] * BALL_RADIUS, Q[1] + N[1] * BALL_RADIUS
                    collided = True
                    break

    # Handle ball-ball collisions
    for i in range(len(balls)):
        for j in range(i + 1, len(balls)):
            handle_ball_collision(balls[i], balls[j])

    # Render
    screen.fill((0, 0, 0))
    pygame.draw.lines(screen, (255, 255, 255), True, vertices, 2)
    for ball in balls:
        pygame.draw.circle(screen, ball.color, (int(ball.x), int(ball.y)), BALL_RADIUS)
    pygame.display.flip()
    clock.tick(60)

pygame.quit()